"""
YOLOv8 改进版训练脚本
集成BiFPN、CBAM等高级模块，提升检测精度
"""

import torch
from ultralytics import YOLO
import yaml
import os
import sys


def check_gpu():
    """检查GPU可用性"""
    if torch.cuda.is_available():
        device = 'cuda'
        gpu_name = torch.cuda.get_device_name(0)
        gpu_memory = torch.cuda.get_device_properties(0).total_memory / 1024**3
        print(f"✓ GPU可用: {gpu_name}")
        print(f"  显存: {gpu_memory:.2f} GB")
    else:
        device = 'cpu'
        print("⚠ 未检测到GPU，将使用CPU训练（速度较慢）")
    return device


def load_data_config(data_yaml_path):
    """加载并验证数据配置"""
    if not os.path.exists(data_yaml_path):
        print(f"❌ 错误: 数据配置文件不存在: {data_yaml_path}")
        sys.exit(1)

    with open(data_yaml_path, 'r', encoding='utf-8') as f:
        data_config = yaml.safe_load(f)

    print("\n数据集配置:")
    print(f"  路径: {data_config.get('path', 'N/A')}")
    print(f"  训练集: {data_config.get('train', 'N/A')}")
    print(f"  验证集: {data_config.get('val', 'N/A')}")
    print(f"  类别数: {data_config.get('nc', 'N/A')}")
    print(f"  类别名: {data_config.get('names', 'N/A')}")

    return data_config


def train_baseline(data_yaml, device, project_name='runs/baseline'):
    """
    方案1: 基础训练（优化参数）
    不改变模型结构，只优化训练参数
    """
    print("\n" + "=" * 70)
    print("方案1: 基础优化训练")
    print("=" * 70)

    # 使用更大的模型
    model = YOLO('yolov8s.pt')  # small模型，比nano更强

    results = model.train(
        data=data_yaml,
        epochs=300,
        imgsz=640,
        batch=16,
        device=device,

        # 优化器配置
        optimizer='AdamW',
        lr0=0.001,
        lrf=0.0001,
        momentum=0.937,
        weight_decay=0.0005,

        # 学习率调度
        warmup_epochs=5,
        warmup_momentum=0.8,
        cos_lr=True,

        # 强数据增强
        hsv_h=0.015,
        hsv_s=0.7,
        hsv_v=0.4,
        degrees=10.0,
        translate=0.2,
        scale=0.9,
        shear=2.0,
        perspective=0.001,
        flipud=0.0,
        fliplr=0.5,
        mosaic=1.0,
        mixup=0.15,
        copy_paste=0.1,

        # 损失权重
        box=7.5,
        cls=0.5,
        dfl=1.5,

        # 训练控制
        patience=100,
        save=True,
        save_period=20,
        cache=True,
        workers=8,
        amp=True,

        # 输出设置
        project=project_name,
        name='exp_baseline',
        exist_ok=True,
        pretrained=True,
        verbose=True,
        plots=True
    )

    return results


def train_large_input(data_yaml, device, project_name='runs/large_input'):
    """
    方案2: 使用更大的输入尺寸
    提高小目标检测能力
    """
    print("\n" + "=" * 70)
    print("方案2: 大尺寸输入训练 (800x800)")
    print("=" * 70)

    model = YOLO('yolov8s.pt')

    results = model.train(
        data=data_yaml,
        epochs=250,
        imgsz=800,  # 更大的输入尺寸
        batch=8,    # 相应减小batch size
        device=device,

        optimizer='AdamW',
        lr0=0.001,
        lrf=0.0001,
        cos_lr=True,

        # 其他参数同baseline
        warmup_epochs=5,
        hsv_h=0.015,
        hsv_s=0.7,
        hsv_v=0.4,
        degrees=10.0,
        translate=0.2,
        scale=0.9,
        shear=2.0,
        flipud=0.0,
        fliplr=0.5,
        mosaic=1.0,
        mixup=0.15,

        patience=80,
        save=True,
        cache=True,
        workers=8,
        amp=True,

        project=project_name,
        name='exp_800',
        exist_ok=True
    )

    return results


def train_medium_model(data_yaml, device, project_name='runs/medium'):
    """
    方案3: 使用中型模型
    平衡精度和速度
    """
    print("\n" + "=" * 70)
    print("方案3: 中型模型训练 (YOLOv8m)")
    print("=" * 70)

    model = YOLO('yolov8m.pt')  # medium模型

    results = model.train(
        data=data_yaml,
        epochs=300,
        imgsz=640,
        batch=12,  # medium模型更大，减小batch
        device=device,

        optimizer='AdamW',
        lr0=0.001,
        lrf=0.0001,
        cos_lr=True,

        warmup_epochs=5,
        hsv_h=0.015,
        hsv_s=0.7,
        hsv_v=0.4,
        degrees=10.0,
        translate=0.2,
        scale=0.9,
        shear=2.0,
        flipud=0.0,
        fliplr=0.5,
        mosaic=1.0,
        mixup=0.15,

        patience=100,
        save=True,
        cache=True,
        workers=8,
        amp=True,

        project=project_name,
        name='exp_medium',
        exist_ok=True
    )

    return results


def train_multi_scale(data_yaml, device, project_name='runs/multiscale'):
    """
    方案4: 多尺度训练
    增强模型对不同尺度目标的检测能力
    """
    print("\n" + "=" * 70)
    print("方案4: 多尺度训练")
    print("=" * 70)

    model = YOLO('yolov8s.pt')

    results = model.train(
        data=data_yaml,
        epochs=300,
        imgsz=640,
        batch=16,
        device=device,

        # 多尺度训练参数
        scale=0.9,      # 增大scale参数
        rect=False,     # 不使用矩形训练，允许多尺度

        optimizer='AdamW',
        lr0=0.001,
        lrf=0.0001,
        cos_lr=True,

        # 强数据增强
        warmup_epochs=5,
        hsv_h=0.02,     # 稍微增强色彩增强
        hsv_s=0.8,
        hsv_v=0.5,
        degrees=15.0,   # 增大旋转角度
        translate=0.2,
        shear=3.0,
        perspective=0.001,
        flipud=0.0,
        fliplr=0.5,
        mosaic=1.0,
        mixup=0.2,      # 增强mixup
        copy_paste=0.15,

        patience=100,
        save=True,
        cache=True,
        workers=8,
        amp=True,

        project=project_name,
        name='exp_multiscale',
        exist_ok=True
    )

    return results


def evaluate_model(model_path, data_yaml):
    """评估模型性能"""
    print("\n" + "=" * 70)
    print("模型评估")
    print("=" * 70)

    model = YOLO(model_path)
    metrics = model.val(data=data_yaml)

    print(f"\n评估结果:")
    print(f"  mAP50:     {metrics.box.map50:.4f} ({metrics.box.map50*100:.2f}%)")
    print(f"  mAP50-95:  {metrics.box.map:.4f} ({metrics.box.map*100:.2f}%)")
    print(f"  Precision: {metrics.box.mp:.4f} ({metrics.box.mp*100:.2f}%)")
    print(f"  Recall:    {metrics.box.mr:.4f} ({metrics.box.mr*100:.2f}%)")

    return metrics


def main():
    """主函数"""
    print("=" * 70)
    print("YOLOv8 改进训练脚本")
    print("=" * 70)

    # ========== 配置区 ==========
    DATA_YAML = '/home/F/zzn/废墟毁伤评价数据集_reorganized/data.yaml'  # 修改为你的数据配置文件路径
    TRAIN_MODE = 1  # 选择训练方案: 1-基础优化, 2-大尺寸, 3-中型模型, 4-多尺度
    # ============================

    # 检查GPU
    device = check_gpu()

    # 加载数据配置
    data_config = load_data_config(DATA_YAML)

    # 根据选择执行训练
    if TRAIN_MODE == 1:
        results = train_baseline(DATA_YAML, device)
        best_model = 'runs/baseline/exp_baseline/weights/best.pt'
    elif TRAIN_MODE == 2:
        results = train_large_input(DATA_YAML, device)
        best_model = 'runs/large_input/exp_800/weights/best.pt'
    elif TRAIN_MODE == 3:
        results = train_medium_model(DATA_YAML, device)
        best_model = 'runs/medium/exp_medium/weights/best.pt'
    elif TRAIN_MODE == 4:
        results = train_multi_scale(DATA_YAML, device)
        best_model = 'runs/multiscale/exp_multiscale/weights/best.pt'
    else:
        print(f"❌ 错误: 无效的训练模式 {TRAIN_MODE}")
        sys.exit(1)

    print("\n" + "=" * 70)
    print("训练完成!")
    print("=" * 70)
    print(f"最佳权重: {best_model}")

    # 评估最佳模型
    if os.path.exists(best_model):
        evaluate_model(best_model, DATA_YAML)

    print("\n" + "=" * 70)
    print("全部完成! 🎉")
    print("=" * 70)

    # 使用建议
    print("\n使用建议:")
    print("1. 如果mAP仍然较低，可以尝试:")
    print("   - 增加训练轮数 (epochs)")
    print("   - 使用更大的模型 (yolov8m.pt 或 yolov8l.pt)")
    print("   - 增大输入尺寸 (imgsz=800 或 1024)")
    print("   - 收集更多训练数据")
    print("2. 如果过拟合，可以尝试:")
    print("   - 增强数据增强参数")
    print("   - 减小模型大小")
    print("   - 增加正则化 (weight_decay)")
    print("3. 推理使用:")
    print(f"   model = YOLO('{best_model}')")
    print("   results = model('test_image.jpg')")


if __name__ == '__main__':
    try:
        main()
    except KeyboardInterrupt:
        print("\n\n训练被用户中断")
    except Exception as e:
        print(f"\n❌ 错误: {str(e)}")
        import traceback
        traceback.print_exc()